Implement tree (prefix tree)

Time: O(N); Space: O(1); medium

Implement a Trie with:

  • insert()

  • search()

  • startsWith()

methods.

Notes:

  • Time O(N), per operation

  • You may assume that all inputs are consist of lowercase letters a-z.

[1]:

class TrieNode(object):
    # Initialize your data structure here
    def __init__(self):
        self.is_string = False
        self.leaves = {}

class Trie(object):

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        '''
        Inserts a word into the Trie
        :type word: str
        :rtype: none
        '''
        cur = self.root
        for c in word:
            if not c in cur.leaves:
                cur.leaves[c] = TrieNode()
            cur = cur.leaves[c]
        cur.is_string = True

    def search(self, word) -> bool:
        '''
        Returns if the word is in the Trie
        :type word: str
        :rtype: bool
        '''
        node = self.childSearch(word)
        if node:
            return node.is_string
        return False

    def startsWith(self, prefix) -> bool:
        '''
        Returns if there is any word in the trie that starts with the given prefix.
        :type prefix: str
        :rtype: bool
        '''
        return self.childSearch(prefix) is not None

    def childSearch(self, word):
        cur = self.root
        for c in word:
            if c in cur.leaves:
                cur = cur.leaves[c]
            else:
                return None
        return cur

[2]:

# Your Trie object will be instantiated and called as such:
trie = Trie()
trie.insert("somestring")
assert trie.search("key") == False
trie.insert("abc")
assert trie.search("abc") == True

trie.insert("lintcode")
assert trie.search("code") == False

assert trie.startsWith("lint") == True
assert trie.startsWith("linterror") == False

trie.insert("linterror")
assert trie.search("lintcode") == True
assert trie.startsWith("linterror") == True